{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget https://github.com/aisingapore/seacorenlp-data/raw/main/id/constituency/train.txt\n",
    "# !wget https://github.com/aisingapore/seacorenlp-data/raw/main/id/constituency/test.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-10-06 13:14:50.801742: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-10-06 13:14:50.875726: I tensorflow/core/util/port.cc:104] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2023-10-06 13:14:51.315980: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory\n",
      "2023-10-06 13:14:51.316039: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory\n",
      "2023-10-06 13:14:51.316042: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n",
      "/home/husein/dev/malaya/malaya/tokenizer.py:214: FutureWarning: Possible nested set at position 3397\n",
      "  self.tok = re.compile(r'({})'.format('|'.join(pipeline)))\n",
      "/home/husein/dev/malaya/malaya/tokenizer.py:214: FutureWarning: Possible nested set at position 3927\n",
      "  self.tok = re.compile(r'({})'.format('|'.join(pipeline)))\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from malaya.function.constituency import evaluate, trees_newline as trees\n",
    "from transformers import AutoTokenizer, T5Config\n",
    "from malaya.torch_model.t5 import T5Constituency\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "BERT_TOKEN_MAPPING = {\n",
    "    \"-LRB-\": \"(\",\n",
    "    \"-RRB-\": \")\",\n",
    "    \"-LCB-\": \"{\",\n",
    "    \"-RCB-\": \"}\",\n",
    "    \"-LSB-\": \"[\",\n",
    "    \"-RSB-\": \"]\",\n",
    "    \"``\": '\"',\n",
    "    \"''\": '\"',\n",
    "    \"`\": \"'\",\n",
    "    '«': '\"',\n",
    "    '»': '\"',\n",
    "    '‘': \"'\",\n",
    "    '’': \"'\",\n",
    "    '“': '\"',\n",
    "    '”': '\"',\n",
    "    '„': '\"',\n",
    "    '‹': \"'\",\n",
    "    '›': \"'\",\n",
    "    \"\\u2013\": \"--\",\n",
    "    \"\\u2014\": \"--\",\n",
    "    }\n",
    "\n",
    "def process_word(word):\n",
    "    word = word.replace('\\\\/', '/').replace('\\\\*', '*')\n",
    "    word = word.replace('-LSB-', '[').replace('-RSB-', ']')\n",
    "    word = word.replace('-LRB-', '(').replace('-RRB-', ')')\n",
    "    if word == \"n't\" and cleaned_words:\n",
    "        cleaned_words[-1] = cleaned_words[-1] + \"n\"\n",
    "        word = \"'t\"\n",
    "    return word"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "\n",
    "class Vocabulary(object):\n",
    "    def __init__(self):\n",
    "        self.frozen = False\n",
    "        self.values = []\n",
    "        self.indices = {}\n",
    "        self.counts = collections.defaultdict(int)\n",
    "\n",
    "    @property\n",
    "    def size(self):\n",
    "        return len(self.values)\n",
    "\n",
    "    def value(self, index):\n",
    "        assert 0 <= index < len(self.values)\n",
    "        return self.values[index]\n",
    "\n",
    "    def index(self, value):\n",
    "        if not self.frozen:\n",
    "            self.counts[value] += 1\n",
    "\n",
    "        if value in self.indices:\n",
    "            return self.indices[value]\n",
    "\n",
    "        elif not self.frozen:\n",
    "            self.values.append(value)\n",
    "            self.indices[value] = len(self.values) - 1\n",
    "            return self.indices[value]\n",
    "\n",
    "        else:\n",
    "            raise ValueError(\"Unknown value: {}\".format(value))\n",
    "\n",
    "    def index_or_unk(self, value, unk_value):\n",
    "        assert self.frozen\n",
    "        if value in self.indices:\n",
    "            return self.indices[value]\n",
    "        else:\n",
    "            return self.indices[unk_value]\n",
    "\n",
    "    def count(self, value):\n",
    "        return self.counts[value]\n",
    "\n",
    "    def freeze(self):\n",
    "        self.frozen = True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_treebank = trees.load_trees('train.txt')\n",
    "train_parse = [tree.convert() for tree in train_treebank]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_treebank = trees.load_trees('test.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e208ebfd1fe498d83753b7dc60dd123",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)okenizer_config.json:   0%|          | 0.00/1.86k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ac3407e259f4cb094b8d19188b2c43e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading spiece.model:   0%|          | 0.00/803k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ae903ba40ada42dca549e720c80d5fdb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/1.79k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading the tokenizer from the `special_tokens_map.json` and the `added_tokens.json` will be removed in `transformers 5`,  it is kept for forward compatibility, but it is recommended to update your `tokenizer_config.json` by uploading it again. You will see the new `added_tokens_decoder` attribute that will store the relevant information.\n",
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. If you see this, DO NOT PANIC! This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thouroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('mesolitica/t5-base-standard-bahasa-cased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_vocab = Vocabulary()\n",
    "label_vocab.index(())\n",
    "\n",
    "tag_vocab = Vocabulary()\n",
    "START = '<s>'\n",
    "STOP = '</s>'\n",
    "UNK = tokenizer.unk_token\n",
    "TAG_UNK = \"UNK\"\n",
    "tag_vocab.index(START)\n",
    "tag_vocab.index(STOP)\n",
    "tag_vocab.index(TAG_UNK)\n",
    "\n",
    "for tree in train_parse:\n",
    "    nodes = [tree]\n",
    "    while nodes:\n",
    "        node = nodes.pop()\n",
    "        if isinstance(node, trees.InternalParseNode):\n",
    "            label_vocab.index(node.label)\n",
    "            nodes.extend(reversed(node.children))\n",
    "        else:\n",
    "            tag_vocab.index(node.tag)\n",
    "            \n",
    "tag_vocab.freeze()\n",
    "label_vocab.freeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "243564db039444029a379bf0c1d3e080",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)lve/main/config.json:   0%|          | 0.00/670 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "config = T5Config.from_pretrained('mesolitica/t5-base-standard-bahasa-cased')\n",
    "config.num_labels = label_vocab.size\n",
    "config.num_tags = tag_vocab.size\n",
    "config.tag_loss_scale = 1.0\n",
    "config.label_vocab = {str(k): v for k, v in label_vocab.indices.items()}\n",
    "config.tag_vocab = tag_vocab.indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9450f0e5da9d4eb79596f0df1dccf6ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading pytorch_model.bin:   0%|          | 0.00/892M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of T5Constituency were not initialized from the model checkpoint at mesolitica/t5-base-standard-bahasa-cased and are newly initialized: ['encoder_encoder.attn_2.proj1.weight', 'encoder_encoder.attn_0.layer_norm.a_2', 'encoder_encoder.attn_1.w_vs2', 'encoder_encoder.ff_1.layer_norm.b_2', 'encoder_encoder.attn_5.layer_norm.a_2', 'encoder_encoder.attn_6.w_ks1', 'encoder_encoder.ff_0.w_2c.weight', 'encoder_encoder.attn_0.w_qs2', 'encoder_encoder.attn_7.w_ks2', 'encoder_encoder.ff_2.w_1c.bias', 'encoder_encoder.ff_7.w_1c.weight', 'encoder_encoder.attn_3.proj2.weight', 'encoder_encoder.ff_0.w_2c.bias', 'encoder_encoder.ff_2.layer_norm.a_2', 'encoder_encoder.ff_3.w_2p.weight', 'encoder_encoder.attn_2.layer_norm.a_2', 'encoder_encoder.ff_1.w_1p.weight', 'encoder_encoder.attn_0.layer_norm.b_2', 'f_label.1.a_2', 'encoder_encoder.ff_1.w_2p.weight', 'encoder_encoder.ff_5.w_1p.bias', 'encoder_encoder.ff_4.w_1p.weight', 'encoder_encoder.ff_0.w_1c.weight', 'encoder_encoder.ff_5.w_2p.bias', 'encoder_encoder.ff_2.w_2c.weight', 'encoder_encoder.attn_7.proj1.weight', 'encoder_encoder.attn_5.w_ks1', 'encoder_encoder.attn_2.w_vs1', 'encoder_encoder.attn_5.w_vs1', 'encoder_encoder.ff_5.w_2c.bias', 'f_tag.0.weight', 'encoder_encoder.attn_1.proj1.weight', 'encoder_encoder.ff_6.w_2c.bias', 'f_tag.1.a_2', 'encoder_encoder.ff_6.w_2p.weight', 'embedding.position_table', 'encoder_encoder.ff_5.layer_norm.b_2', 'encoder_encoder.attn_1.layer_norm.b_2', 'encoder_encoder.ff_5.layer_norm.a_2', 'encoder_encoder.attn_3.w_ks1', 'encoder_encoder.attn_5.proj2.weight', 'f_tag.1.b_2', 'encoder_encoder.attn_6.proj2.weight', 'encoder_encoder.ff_1.w_1c.weight', 'encoder_encoder.attn_0.proj1.weight', 'encoder_encoder.ff_6.w_1p.weight', 'encoder_encoder.ff_2.w_2p.weight', 'encoder_encoder.attn_0.w_ks1', 'encoder_encoder.attn_3.w_vs1', 'encoder_encoder.ff_1.layer_norm.a_2', 'encoder_encoder.attn_3.w_vs2', 'f_label.3.weight', 'encoder_encoder.attn_0.w_qs1', 'encoder_encoder.attn_5.w_qs2', 'encoder_encoder.attn_6.w_vs1', 'encoder_encoder.attn_6.layer_norm.b_2', 'encoder_encoder.attn_2.w_qs1', 'encoder_encoder.ff_7.w_2c.bias', 'encoder_encoder.attn_4.w_ks1', 'encoder_encoder.attn_6.w_ks2', 'encoder_encoder.ff_7.w_1p.bias', 'encoder_encoder.attn_0.w_vs2', 'project_bert.weight', 'encoder_encoder.attn_1.w_ks2', 'encoder_encoder.ff_3.w_2c.weight', 'f_tag.3.bias', 'encoder_encoder.ff_2.w_2c.bias', 'encoder_encoder.ff_7.w_1p.weight', 'encoder_encoder.ff_1.w_2p.bias', 'encoder_encoder.ff_2.w_2p.bias', 'f_label.0.bias', 'encoder_encoder.ff_3.w_2c.bias', 'embedding.layer_norm.a_2', 'encoder_encoder.attn_3.w_qs1', 'encoder_encoder.ff_4.w_1p.bias', 'encoder_encoder.ff_0.w_2p.bias', 'encoder_encoder.ff_4.w_2p.bias', 'encoder_encoder.ff_6.w_1c.weight', 'encoder_encoder.attn_7.proj2.weight', 'encoder_encoder.ff_1.w_2c.bias', 'encoder_encoder.attn_3.layer_norm.b_2', 'encoder_encoder.attn_1.w_qs1', 'encoder_encoder.attn_4.w_vs2', 'encoder_encoder.attn_4.w_qs1', 'encoder_encoder.attn_5.w_vs2', 'encoder_encoder.ff_7.layer_norm.a_2', 'encoder_encoder.attn_6.layer_norm.a_2', 'encoder_encoder.ff_0.w_1p.weight', 'encoder_encoder.ff_6.w_2c.weight', 'encoder_encoder.attn_4.w_vs1', 'encoder_encoder.attn_2.w_vs2', 'encoder_encoder.attn_6.proj1.weight', 'encoder_encoder.ff_7.w_2p.weight', 'encoder_encoder.attn_1.w_vs1', 'encoder_encoder.attn_2.layer_norm.b_2', 'encoder_encoder.ff_4.w_2p.weight', 'encoder_encoder.attn_4.layer_norm.b_2', 'encoder_encoder.ff_4.layer_norm.b_2', 'encoder_encoder.ff_5.w_1c.weight', 'encoder_encoder.ff_5.w_2p.weight', 'encoder_encoder.attn_6.w_qs1', 'encoder_encoder.ff_6.w_1p.bias', 'encoder_encoder.attn_7.w_qs2', 'encoder_encoder.ff_7.w_2c.weight', 'encoder_encoder.attn_4.layer_norm.a_2', 'f_tag.0.bias', 'encoder_encoder.attn_7.layer_norm.b_2', 'encoder_encoder.ff_6.layer_norm.a_2', 'encoder_encoder.attn_1.w_qs2', 'encoder_encoder.ff_5.w_2c.weight', 'encoder_encoder.ff_1.w_1p.bias', 'encoder_encoder.ff_4.w_1c.bias', 'encoder_encoder.ff_3.layer_norm.b_2', 'encoder_encoder.ff_0.w_1c.bias', 'encoder_encoder.ff_2.w_1p.bias', 'encoder_encoder.ff_3.layer_norm.a_2', 'encoder_encoder.ff_6.w_1c.bias', 'encoder_encoder.attn_3.w_ks2', 'encoder_encoder.attn_0.w_ks2', 'encoder_encoder.attn_0.w_vs1', 'encoder_encoder.ff_4.w_2c.bias', 'encoder_encoder.ff_0.layer_norm.a_2', 'encoder_encoder.attn_5.proj1.weight', 'encoder_encoder.attn_7.w_qs1', 'encoder_encoder.ff_2.w_1c.weight', 'encoder_encoder.ff_7.w_1c.bias', 'encoder_encoder.attn_7.layer_norm.a_2', 'encoder_encoder.attn_2.w_qs2', 'f_label.0.weight', 'encoder_encoder.attn_3.layer_norm.a_2', 'encoder_encoder.ff_1.w_1c.bias', 'encoder_encoder.attn_1.proj2.weight', 'encoder_encoder.ff_5.w_1p.weight', 'encoder_encoder.attn_4.w_ks2', 'encoder_encoder.ff_4.w_2c.weight', 'encoder_encoder.ff_6.w_2p.bias', 'encoder_encoder.ff_6.layer_norm.b_2', 'encoder_encoder.attn_7.w_vs2', 'encoder_encoder.attn_3.proj1.weight', 'encoder_encoder.ff_3.w_1p.bias', 'encoder_encoder.attn_1.layer_norm.a_2', 'encoder_encoder.attn_4.proj1.weight', 'encoder_encoder.attn_7.w_ks1', 'encoder_encoder.attn_3.w_qs2', 'encoder_encoder.ff_5.w_1c.bias', 'encoder_encoder.ff_3.w_1p.weight', 'encoder_encoder.attn_1.w_ks1', 'encoder_encoder.attn_5.w_qs1', 'encoder_encoder.attn_2.proj2.weight', 'f_label.3.bias', 'encoder_encoder.ff_1.w_2c.weight', 'encoder_encoder.attn_4.proj2.weight', 'encoder_encoder.ff_0.w_1p.bias', 'embedding.layer_norm.b_2', 'encoder_encoder.ff_3.w_1c.weight', 'encoder_encoder.ff_7.w_2p.bias', 'encoder_encoder.ff_7.layer_norm.b_2', 'encoder_encoder.attn_7.w_vs1', 'encoder_encoder.attn_6.w_vs2', 'encoder_encoder.attn_5.w_ks2', 'f_label.1.b_2', 'encoder_encoder.ff_0.layer_norm.b_2', 'encoder_encoder.attn_2.w_ks2', 'f_tag.3.weight', 'encoder_encoder.ff_2.layer_norm.b_2', 'encoder_encoder.attn_0.proj2.weight', 'encoder_encoder.ff_0.w_2p.weight', 'encoder_encoder.attn_5.layer_norm.b_2', 'encoder_encoder.ff_2.w_1p.weight', 'encoder_encoder.ff_3.w_1c.bias', 'encoder_encoder.ff_4.layer_norm.a_2', 'encoder_encoder.ff_3.w_2p.bias', 'encoder_encoder.attn_2.w_ks1', 'encoder_encoder.ff_4.w_1c.weight', 'encoder_encoder.attn_6.w_qs2', 'encoder_encoder.attn_4.w_qs2']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "model = T5Constituency.from_pretrained('mesolitica/t5-base-standard-bahasa-cased', config = config)\n",
    "_ = model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainable_parameters = [param for param in model.parameters() if param.requires_grad]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = torch.optim.AdamW(trainable_parameters, lr = 2e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BatchIndices:\n",
    "    def __init__(self, batch_idxs_np):\n",
    "        self.batch_idxs_np = batch_idxs_np\n",
    "        self.batch_idxs_torch = torch.from_numpy(batch_idxs_np)\n",
    "        self.batch_size = int(1 + np.max(batch_idxs_np))\n",
    "\n",
    "        batch_idxs_np_extra = np.concatenate([[-1], batch_idxs_np, [-1]])\n",
    "        self.boundaries_np = np.nonzero(batch_idxs_np_extra[1:] != batch_idxs_np_extra[:-1])[0]\n",
    "        self.seq_lens_np = self.boundaries_np[1:] - self.boundaries_np[:-1]\n",
    "        assert len(self.seq_lens_np) == self.batch_size\n",
    "        self.max_len = int(np.max(self.boundaries_np[1:] - self.boundaries_np[:-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_batch(sentences, golds, subbatch_max_tokens=3000):\n",
    "    lens = [\n",
    "        len(tokenizer.tokenize(' '.join([word for (_, word) in sentence]))) + 2\n",
    "        for sentence in sentences\n",
    "    ]\n",
    "\n",
    "    lens = np.asarray(lens, dtype=int)\n",
    "    lens_argsort = np.argsort(lens).tolist()\n",
    "\n",
    "    num_subbatches = 0\n",
    "    subbatch_size = 1\n",
    "    while lens_argsort:\n",
    "        if (subbatch_size == len(lens_argsort)) or (subbatch_size * lens[lens_argsort[subbatch_size]] > subbatch_max_tokens):\n",
    "            yield [sentences[i] for i in lens_argsort[:subbatch_size]], [golds[i] for i in lens_argsort[:subbatch_size]]\n",
    "            lens_argsort = lens_argsort[subbatch_size:]\n",
    "            num_subbatches += 1\n",
    "            subbatch_size = 1\n",
    "        else:\n",
    "            subbatch_size += 1\n",
    "            \n",
    "def pad_sentence_batch(sentence_batch, pad_int):\n",
    "    padded_seqs = []\n",
    "    seq_lens = []\n",
    "    max_sentence_len = max([len(sentence) for sentence in sentence_batch])\n",
    "    for sentence in sentence_batch:\n",
    "        padded_seqs.append(\n",
    "            sentence + [pad_int] * (max_sentence_len - len(sentence))\n",
    "        )\n",
    "        seq_lens.append(len(sentence))\n",
    "    return padded_seqs, seq_lens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process(sentences, golds = None):\n",
    "    \n",
    "    all_input_ids = []\n",
    "    all_word_start_mask = []\n",
    "    all_word_end_mask = []\n",
    "\n",
    "    for snum, sentence in enumerate(sentences):\n",
    "\n",
    "        tokens = []\n",
    "        word_start_mask = []\n",
    "        word_end_mask = []\n",
    "        tokens.append(START)\n",
    "        word_start_mask.append(1)\n",
    "        word_end_mask.append(1)\n",
    "\n",
    "        cleaned_words = []\n",
    "        for _, word in sentence:\n",
    "            cleaned_words.append(process_word(word))\n",
    "\n",
    "        for word in cleaned_words:\n",
    "            word_tokens = tokenizer.tokenize(word)\n",
    "            for _ in range(len(word_tokens)):\n",
    "                word_start_mask.append(0)\n",
    "                word_end_mask.append(0)\n",
    "            word_start_mask[len(tokens)] = 1\n",
    "            word_end_mask[-1] = 1\n",
    "            tokens.extend(word_tokens)\n",
    "        tokens.append(STOP)\n",
    "        word_start_mask.append(1)\n",
    "        word_end_mask.append(1)\n",
    "\n",
    "        input_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "        all_input_ids.append(input_ids)\n",
    "        all_word_start_mask.append(word_start_mask)\n",
    "        all_word_end_mask.append(word_end_mask)\n",
    "    \n",
    "    padded = tokenizer.pad({\n",
    "        'input_ids': all_input_ids,\n",
    "    }, return_tensors = 'pt')\n",
    "\n",
    "    all_word_start_mask = torch.from_numpy(np.array(pad_sentence_batch(all_word_start_mask, 0)[0]))\n",
    "    all_word_end_mask = torch.from_numpy(np.array(pad_sentence_batch(all_word_end_mask, 0)[0]))\n",
    "    \n",
    "    padded['sentences'] = sentences\n",
    "    padded['all_word_start_mask'] = all_word_start_mask\n",
    "    padded['all_word_end_mask'] = all_word_end_mask\n",
    "    \n",
    "    packed_len = sum([(len(sentence) + 2) for sentence in sentences])\n",
    "    i = 0\n",
    "    tag_idxs = np.zeros(packed_len, dtype=int)\n",
    "    batch_idxs = np.zeros(packed_len, dtype=int)\n",
    "    for snum, sentence in enumerate(sentences):\n",
    "        for (tag, word) in [(START, START)] + sentence + [(STOP, STOP)]:\n",
    "            if golds is not None:\n",
    "                tag_idxs[i] = tag_vocab.index_or_unk(tag, TAG_UNK)\n",
    "            else:\n",
    "                tag_idxs[i] = 0\n",
    "            batch_idxs[i] = snum\n",
    "            i += 1\n",
    "    \n",
    "    batch_idxs = BatchIndices(batch_idxs)\n",
    "    padded['batch_idxs'] = batch_idxs\n",
    "    tag_idxs = torch.from_numpy(tag_idxs)\n",
    "    padded['tag_idxs'] = tag_idxs\n",
    "    \n",
    "    if golds is not None:\n",
    "        gold_tag_idxs = tag_idxs\n",
    "        padded['gold_tag_idxs'] = gold_tag_idxs\n",
    "        padded['golds'] = golds\n",
    "        \n",
    "        \n",
    "    if torch.cuda.is_available():\n",
    "        for k in padded.keys():\n",
    "            if isinstance(padded[k], torch.Tensor):\n",
    "                padded[k] = padded[k].cuda()\n",
    "    \n",
    "        padded['batch_idxs'].batch_idxs_torch = padded['batch_idxs'].batch_idxs_torch.cuda()\n",
    "    \n",
    "    return padded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a T5TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor(264.3695, device='cuda:0', grad_fn=<AddBackward0>), tensor(277.7802, device='cuda:0', grad_fn=<MulBackward0>))\n"
     ]
    }
   ],
   "source": [
    "start_index = 0\n",
    "batch_size = 4\n",
    "batch_trees = train_parse[start_index:start_index + batch_size]\n",
    "batch_sentences = [[(leaf.tag, leaf.word) for leaf in tree.leaves()] for tree in batch_trees]\n",
    "for subbatch_sentences, subbatch_trees in split_batch(batch_sentences, batch_trees):\n",
    "    print(model(**process(subbatch_sentences, subbatch_trees)))\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# model(**process(subbatch_sentences)), model(**process(subbatch_sentences, subbatch_trees))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:03<00:00,  7.92it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0, loss: 7.5798946719169615, dev_fscore: (Recall=69.07, Precision=71.65, FScore=70.34, CompleteMatch=7.80, TaggingAccuracy=93.90)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:00<00:00,  8.21it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1, loss: 3.7848024014234545, dev_fscore: (Recall=71.99, Precision=75.98, FScore=73.93, CompleteMatch=11.10, TaggingAccuracy=94.22)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:01<00:00,  8.11it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 2, loss: 2.6138466329574586, dev_fscore: (Recall=74.78, Precision=77.75, FScore=76.24, CompleteMatch=14.40, TaggingAccuracy=94.58)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:01<00:00,  8.09it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 3, loss: 1.8761003473997115, dev_fscore: (Recall=77.29, Precision=78.71, FScore=78.00, CompleteMatch=15.90, TaggingAccuracy=94.20)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:03<00:00,  7.92it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 4, loss: 1.4609425439834596, dev_fscore: (Recall=78.53, Precision=79.50, FScore=79.01, CompleteMatch=16.70, TaggingAccuracy=94.64)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:02<00:00,  7.99it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 5, loss: 1.176018776535988, dev_fscore: (Recall=78.50, Precision=80.40, FScore=79.44, CompleteMatch=18.90, TaggingAccuracy=94.46)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:02<00:00,  8.05it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 6, loss: 1.0093939388394355, dev_fscore: (Recall=80.66, Precision=80.04, FScore=80.35, CompleteMatch=19.90, TaggingAccuracy=94.84)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:02<00:00,  8.00it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 7, loss: 0.8886480987668037, dev_fscore: (Recall=80.30, Precision=80.29, FScore=80.30, CompleteMatch=20.30, TaggingAccuracy=94.80)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:02<00:00,  8.00it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 8, loss: 0.7949668666422367, dev_fscore: (Recall=80.57, Precision=80.79, FScore=80.68, CompleteMatch=18.90, TaggingAccuracy=94.83)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:02<00:00,  7.96it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 9, loss: 0.719270875453949, dev_fscore: (Recall=81.67, Precision=80.25, FScore=80.96, CompleteMatch=21.30, TaggingAccuracy=94.83)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:01<00:00,  8.08it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 10, loss: 0.6778207948207855, dev_fscore: (Recall=81.26, Precision=80.26, FScore=80.76, CompleteMatch=19.90, TaggingAccuracy=94.93)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:01<00:00,  8.13it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 11, loss: 0.6211680134832859, dev_fscore: (Recall=80.91, Precision=80.50, FScore=80.71, CompleteMatch=21.90, TaggingAccuracy=94.80)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:00<00:00,  8.21it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 12, loss: 0.5911222234815359, dev_fscore: (Recall=82.31, Precision=79.45, FScore=80.85, CompleteMatch=21.10, TaggingAccuracy=94.69)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:00<00:00,  8.22it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 13, loss: 0.5502841906249523, dev_fscore: (Recall=80.79, Precision=79.01, FScore=79.89, CompleteMatch=20.20, TaggingAccuracy=94.65)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:00<00:00,  8.22it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 14, loss: 0.5245251276642084, dev_fscore: (Recall=81.43, Precision=81.11, FScore=81.27, CompleteMatch=21.40, TaggingAccuracy=94.39)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:01<00:00,  8.12it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 15, loss: 0.5087404198795556, dev_fscore: (Recall=82.02, Precision=81.13, FScore=81.57, CompleteMatch=22.50, TaggingAccuracy=94.87)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:01<00:00,  8.15it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 16, loss: 0.4565132545977831, dev_fscore: (Recall=81.30, Precision=81.18, FScore=81.24, CompleteMatch=21.30, TaggingAccuracy=94.87)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:01<00:00,  8.18it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 17, loss: 0.46172504172474144, dev_fscore: (Recall=81.99, Precision=80.96, FScore=81.48, CompleteMatch=21.90, TaggingAccuracy=94.88)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:01<00:00,  8.17it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 18, loss: 0.4225267358124256, dev_fscore: (Recall=81.85, Precision=81.49, FScore=81.67, CompleteMatch=21.30, TaggingAccuracy=94.64)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:01<00:00,  8.18it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 19, loss: 0.42223594155162575, dev_fscore: (Recall=82.23, Precision=82.12, FScore=82.18, CompleteMatch=23.50, TaggingAccuracy=94.69)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:02<00:00,  8.02it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 20, loss: 0.372675468467176, dev_fscore: (Recall=81.92, Precision=81.46, FScore=81.69, CompleteMatch=23.10, TaggingAccuracy=94.74)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:03<00:00,  7.82it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 21, loss: 0.36088745640218256, dev_fscore: (Recall=82.61, Precision=80.66, FScore=81.62, CompleteMatch=22.20, TaggingAccuracy=94.81)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:02<00:00,  7.98it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 22, loss: 0.3422843780517578, dev_fscore: (Recall=82.21, Precision=81.96, FScore=82.09, CompleteMatch=23.40, TaggingAccuracy=95.03)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:02<00:00,  8.01it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 23, loss: 0.3337622888535261, dev_fscore: (Recall=81.49, Precision=81.59, FScore=81.54, CompleteMatch=21.90, TaggingAccuracy=94.96)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:02<00:00,  8.01it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 24, loss: 0.3519724658802152, dev_fscore: (Recall=82.76, Precision=81.06, FScore=81.90, CompleteMatch=22.90, TaggingAccuracy=95.02)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:02<00:00,  8.02it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 25, loss: 0.32064215297996995, dev_fscore: (Recall=81.95, Precision=82.12, FScore=82.04, CompleteMatch=23.60, TaggingAccuracy=94.92)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:02<00:00,  8.04it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 26, loss: 0.27468749369308354, dev_fscore: (Recall=82.52, Precision=81.04, FScore=81.77, CompleteMatch=23.00, TaggingAccuracy=94.88)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:02<00:00,  8.05it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 27, loss: 0.28334507382661106, dev_fscore: (Recall=83.06, Precision=81.21, FScore=82.13, CompleteMatch=23.00, TaggingAccuracy=94.83)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:00<00:00,  8.21it/s]\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 28, loss: 0.27862213914468886, dev_fscore: (Recall=81.60, Precision=82.53, FScore=82.06, CompleteMatch=23.20, TaggingAccuracy=94.70)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 500/500 [01:01<00:00,  8.16it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 29, loss: 0.29495231349393725, dev_fscore: (Recall=80.62, Precision=83.72, FScore=82.14, CompleteMatch=21.50, TaggingAccuracy=94.88)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    }
   ],
   "source": [
    "batch_size = 16\n",
    "epoch = 100\n",
    "\n",
    "best_dev_fscore = -np.inf\n",
    "patient = 10\n",
    "current_patient = 0\n",
    "\n",
    "for e in range(epoch):\n",
    "    pbar = tqdm(range(0, len(train_parse), batch_size))\n",
    "    losses = []\n",
    "    for start_index in pbar:\n",
    "        trainer.zero_grad()\n",
    "        batch_loss_value = 0.0\n",
    "        batch_trees = train_parse[start_index:start_index + batch_size]\n",
    "        batch_sentences = [[(leaf.tag, leaf.word) for leaf in tree.leaves()] for tree in batch_trees]\n",
    "        batch_num_tokens = sum(len(sentence) for sentence in batch_sentences)\n",
    "\n",
    "        for subbatch_sentences, subbatch_trees in split_batch(batch_sentences, batch_trees):\n",
    "            loss, tag_loss =  model(**process(subbatch_sentences, subbatch_trees))\n",
    "            loss = tag_loss / len(subbatch_sentences) + loss / batch_num_tokens\n",
    "            loss_value = float(loss.data.cpu().numpy())\n",
    "            batch_loss_value += loss_value\n",
    "            if loss_value > 0:\n",
    "                loss.backward()\n",
    "        \n",
    "        grad_norm = torch.nn.utils.clip_grad_norm_(trainable_parameters, 5.0)\n",
    "        trainer.step()\n",
    "        losses.append(batch_loss_value)\n",
    "        \n",
    "    dev_predicted = []\n",
    "    for dev_start_index in range(0, len(dev_treebank), batch_size):\n",
    "        subbatch_trees = dev_treebank[dev_start_index:dev_start_index+batch_size]\n",
    "        subbatch_sentences = [[(leaf.tag, leaf.word) for leaf in tree.leaves()] for tree in subbatch_trees]\n",
    "        predicted, _ =  model(**process(subbatch_sentences))\n",
    "        dev_predicted.extend([p.convert() for p in predicted])\n",
    "    \n",
    "    dev_fscore = evaluate.evalb('deprecated/EVALB', dev_treebank, dev_predicted)\n",
    "    \n",
    "    print(f'epoch: {e}, loss: {np.mean(losses)}, dev_fscore: {dev_fscore}')\n",
    "    \n",
    "    if dev_fscore.fscore >= best_dev_fscore:\n",
    "        best_dev_fscore = dev_fscore.fscore\n",
    "        current_patient = 0\n",
    "        model.save_pretrained('base')\n",
    "    else:\n",
    "        current_patient += 1\n",
    "    \n",
    "    if current_patient >= patient:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "82.18"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_dev_fscore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ = T5Constituency.from_pretrained('./base')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "547231d0a85c4ca6be86b81c72f52e3e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/545M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/constituency-parsing-t5-base-standard-bahasa-cased/commit/6c31b09ea74d63447f3403e8a90e18c5c91ca44e', commit_message='Upload T5Constituency', commit_description='', oid='6c31b09ea74d63447f3403e8a90e18c5c91ca44e', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_.push_to_hub('mesolitica/constituency-parsing-t5-base-standard-bahasa-cased', safe_serialization = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "24f9a0a08b604ab6958eb6a20b5adec7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "spiece.model:   0%|          | 0.00/803k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/mesolitica/constituency-parsing-t5-base-standard-bahasa-cased/commit/4806f64f4ab435e2a3923dbc24532e650026236e', commit_message='Upload tokenizer', commit_description='', oid='4806f64f4ab435e2a3923dbc24532e650026236e', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.push_to_hub('mesolitica/constituency-parsing-t5-base-standard-bahasa-cased', safe_serialization = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
